# 写一个python脚本实现以下的功能：
# 1. 输入原始数据集文件夹其中包含images、labels文件夹，要求读取原始数据集中的图片和标签
# 2. 根据原始数据集和设定的train、val、test比例,在指定的outputdir中生成标准yolo数据集。格式是：
# images:
#     -train
#     -val
#     -test
# labels:
#     -train
#     -val
#     -test

import os
import random
import shutil

def read_dataset(dataset_dir):
    """
    读取原始数据集中的图片和标签。
    :param dataset_dir: 原始数据集目录路径
    :return: 图片和标签文件路径列表
    """
    images_dir = os.path.join(dataset_dir, 'images')
    labels_dir = os.path.join(dataset_dir, 'labels')

    image_files = [os.path.join(images_dir, f) for f in os.listdir(images_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    label_files = [os.path.join(labels_dir, f) for f in os.listdir(labels_dir) if f.lower().endswith('.txt')]

    return image_files, label_files


def split_dataset(image_files, label_files, ratios):
    """
    根据比例划分数据集。
    :param image_files: 图片文件路径列表
    :param label_files: 标签文件路径列表
    :param ratios: 划分比例 (train, val, test)
    :return: 划分后的数据集字典
    """
    combined = list(zip(image_files, label_files))
    random.shuffle(combined)

    total_count = len(combined)
    train_count = int(total_count * ratios['train'])
    val_count = int(total_count * ratios['val'])

    train_set = combined[:train_count]
    val_set = combined[train_count:train_count + val_count]
    test_set = combined[train_count + val_count:]

    return {
        'train': train_set,
        'val': val_set,
        'test': test_set
    }


def create_yolo_structure(output_dir):
    """
    创建YOLO格式的数据集目录结构。
    :param output_dir: 输出目录路径
    """
    os.makedirs(os.path.join(output_dir, 'images', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'val'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'test'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'val'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'test'), exist_ok=True)


def copy_files(file_pairs, output_dir, subset):
    """
    复制文件到指定的目标目录。
    :param file_pairs: 图片和标签文件对
    :param output_dir: 输出目录路径
    :param subset: 子集名称 ('train', 'val', 'test')
    """
    for img_path, label_path in file_pairs:
        shutil.copy(img_path, os.path.join(output_dir, 'images', subset, os.path.basename(img_path)))
        shutil.copy(label_path, os.path.join(output_dir, 'labels', subset, os.path.basename(label_path)))


def main():
    # 用户输入配置
    dataset_dir = "D:\Data\minion_detect"
    output_dir = "D:\Project\yolo_dataset_generate\output"
    ratios = {
        'train': 0.85,
        'val': 0.15,
        'test': 0.0
    }

    # 确保比例总和为1
    assert sum(ratios.values()) == 1.0, "比例总和必须为1.0"

    # 读取数据集
    image_files, label_files = read_dataset(dataset_dir)

    # 划分数据集
    splits = split_dataset(image_files, label_files, ratios)

    # 创建YOLO目录结构
    create_yolo_structure(output_dir)

    # 复制文件
    for subset, file_pairs in splits.items():
        copy_files(file_pairs, output_dir, subset)

    print(f"数据集生成完成！")

if __name__ == "__main__":
    main()